#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "riscv.h"
#include "defs.h"
#include "spinlock.h"
#include "trapframe.h"
#include "proc.h"

// S 态内核陷阱入口
extern void kernelvec(void);

// 全局变量：当前trapframe（由kernelvec.S设置）
struct trapframe *kernel_trapframe_ptr = 0;

// ticks
struct spinlock tickslock;
volatile uint64 ticks;

// 简单中断向量表：按 scause 的 irq 编号索引（仅 1/5/9 会用到）
#define MAX_IRQ 64
static interrupt_handler_t irq_table[MAX_IRQ];

// 外部提供的调度钩子；若未提供，链接到弱符号空实现
__attribute__((weak)) void scheduler_tick(void) {}

// IRQ -> sie 位映射（仅 S 态三类）
static inline uint64 sie_bit_for_irq(int irq) {
  switch (irq) {
    case 1:  return SIE_SSIE; // 软件中断
    case 5:  return SIE_STIE; // 时钟中断（未用）
    case 9:  return SIE_SEIE; // 外部中断
    default: return 0;
  }
}

void trapinit(void)
{
  initlock(&tickslock, "time");
  for (int i = 0; i < MAX_IRQ; i++)
    irq_table[i] = 0;

  // 默认注册：软件中断(SSIP, irq=1) → 时钟回调
  register_interrupt(1, timer_interrupt);
}

void trapinithart(void)
{
  // 设置 S 态陷阱向量
  w_stvec((uint64)kernelvec);

  // 仅开启软件中断（SSIP），其他按需开启
  enable_interrupt(1);

  // 注意：不要在启用分页之前打开 S 全局中断
  // 中断将在 kvminithart() 之后由 main() 或其他地方打开
  // w_sstatus(r_sstatus() | SSTATUS_SIE);  // 暂时注释掉
}

void register_interrupt(int irq, interrupt_handler_t h)
{
  if (irq < 0 || irq >= MAX_IRQ) return;
  irq_table[irq] = h;
}

void unregister_interrupt(int irq)
{
  if (irq < 0 || irq >= MAX_IRQ) return;
  irq_table[irq] = 0;
}

void enable_interrupt(int irq)
{
  uint64 bit = sie_bit_for_irq(irq);
  if (bit)
    w_sie(r_sie() | bit);
}

void disable_interrupt(int irq)
{
  uint64 bit = sie_bit_for_irq(irq);
  if (bit)
    w_sie(r_sie() & ~bit);
}

// 设备中断分发：返回 2=时钟路径（SSIP或STIP），1=外设，0=非设备
int devintr(void)
{
  uint64 sc = r_scause();

  if (sc & (1ULL << 63)) {
    // 异步中断
    int irq = (int)(sc & 0xff);

    // S 软件中断（SSIP）：由 M 态 timervec 注入
    if (irq == 1) {
      // 清除 SSIP
      w_sip(r_sip() & ~SIP_SSIP);
      if (irq_table[1])
        irq_table[1]();
      return 2;
    }

    // 外部中断（SEIP，通过 PLIC；此处仅保留钩子）
    if (irq == 9) {
      if (irq_table[9])
        irq_table[9]();
      return 1;
    }

    return 1;
  } else {
    // 同步异常暂不处理
    return 0;
  }
}

void kerneltrap(void)
{
  struct trapframe *tf = kernel_trapframe_ptr;
  
  if (!tf) {
    // 传统处理方式（用于向后兼容）
    uint64 sepc = r_sepc();
    uint64 sstatus = r_sstatus();

    if ((sstatus & SSTATUS_SIE) != 0)
      panic("kerneltrap: interrupts enabled");

    int which = devintr();
    if (which == 0) {
      w_sepc(sepc);
      w_sstatus(sstatus);
      return;
    }

    if (which == 2) {
      struct proc *p = myproc();
      if (p && p->state == RUNNING) {
        yield();
      }
    }

    w_sepc(sepc);
    w_sstatus(sstatus);
    return;
  }

  // 有trapframe，需要区分中断和异常
  // 检查是否是中断（scause最高位为1）
  if (tf->scause & (1ULL << 63)) {
    // 这是中断，调用devintr处理
    int which = devintr();
    if (which != 0) {
      // 如果是时钟中断，进行抢占式调度
      if (which == 2) {
        struct proc *p = myproc();
        if (p && p->state == RUNNING) {
          yield();
        }
      }
      // 中断处理完成，恢复CSR寄存器
      w_sepc(tf->sepc);
      w_sstatus(tf->sstatus);
      kernel_trapframe_ptr = 0; // 清除
      return;
    }
    // 如果不是已知的中断，继续作为异常处理
  }
  
  // 处理同步异常（系统调用或其他异常）
  uint64 cause = tf->scause & ((1ULL<<63)-1);
  
  // 处理系统调用异常（scause = 8/9/11）
  if (cause == 8 || cause == 9 || cause == 11) {
    struct proc *p = myproc();
    if (p) {
      // 设置进程的trapframe
      p->trapframe = tf;
      // 调用系统调用分发器
      syscall_dispatch();
      // 更新sepc（跳过ecall指令）
      tf->sepc += 4;
      // 更新CSR寄存器
      w_sepc(tf->sepc);
      w_sstatus(tf->sstatus);
      kernel_trapframe_ptr = 0; // 清除
      return;
    }
  }
  
  // 其他异常，调用handle_exception
  struct proc *p = myproc();
  handle_exception(tf);
  
  // 检查进程是否被杀死（handle_exception可能设置了killed标志）
  if (p && p->killed) {
    // 进程被杀死，直接退出，不恢复sepc
    exit_process(-1);
    // 不会到达这里
  }
  
  // 只有进程没有被杀死时才恢复sepc
  w_sepc(tf->sepc);
  w_sstatus(tf->sstatus);
  kernel_trapframe_ptr = 0; // 清除
}

// 默认时钟中断处理（在 SSIP 被分发后调用）
void timer_interrupt(void)
{
  acquire(&tickslock);
  ticks++;
  // 唤醒 sleep(&ticks) 的等待者
  wakeup((void*)&ticks);
  release(&tickslock);

  // 心跳钩子（保持轻量）
  scheduler_tick();
}

// 简单的时间读取（CLINT mtime）
uint64 get_time(void)
{
  return *(volatile uint64*)CLINT_MTIME_ADDR;
}